import json
from unicodedata import category
import xml.etree.ElementTree as ET
import sys, os
from tqdm import tqdm



classes = ('aeroplane', 'bicycle', 'bird', 'boat',
        'bottle', 'bus', 'car', 'cat', 'chair',
        'cow', 'diningtable', 'dog', 'horse',
        'motorbike', 'person', 'pottedplant',
        'sheep', 'sofa', 'train', 'tvmonitor')

txtPath = 'E:\dataset\VOC2012\VOCdevkit\VOC2012\ImageSets\Main'
annPath = 'E:\dataset\VOC2012\VOCdevkit\VOC2012\Annotations'
jsonPath = './voc2012_train.json'
groudTruthDict = dict(
    info=[],
    license=[],
    images=[],
    annotations=[],
    categories=[]
)
resList = []




def parse_xml(xmlPath):
    xml = ET.parse(xmlPath)
    fileName = xml.find('filename').text
    width = int(xml.find('size')[0].text)
    height = int(xml.find('size')[1].text)
    objs = xml.findall('object')

    for obj in objs:
        name = obj[0].text
        boxes = [int(obj[4][i].text) for i in range(4)]


def convert(txtPath, annRoot):
    cnt = 1
    with open(txtPath, 'r') as f:
        lines = f.readlines()
        for i, line in enumerate(tqdm(lines)): # 遍历每一张图片的标注
            xmlPath = os.path.join(annRoot, line.strip()+'.xml')
            assert os.path.exists(xmlPath)
            # gts = parse_xml(xmlPath)
            xml = ET.parse(xmlPath)
            fileName = xml.find('filename').text
            width = int(xml.find('size')[0].text)
            height = int(xml.find('size')[1].text)
            groudTruthDict['images'].append(
                {
                    'file_name': fileName,
                    'height': height,
                    'width': width,
                    'id': i
                }
            )
            objs = xml.findall('object')
            for obj in objs:
                name = obj[0].text
                boxes = []
                boxes.append(int(obj.find('bndbox').find('xmin').text))
                boxes.append(int(obj.find('bndbox').find('ymin').text))
                boxes.append(int(obj.find('bndbox').find('xmax').text))
                boxes.append(int(obj.find('bndbox').find('ymax').text))
                
                groudTruthDict['annotations'].append(
                    {
                        'area': width*height,
                        'iscrowd': 0,
                        'image_id': i,
                        'bbox': boxes,
                        'category_id': classes.index(name),
                        'id': cnt
                    }
                )
                cnt += 1
                # resList.
                # append(
                #     {
                #         'image_id': i,
                #         'category_id': classes.index(name),
                #         'bbox': boxes,
                #         'score': 0.99
                #     }
                # )

    for c in classes:
        groudTruthDict['categories'].append(
            {
                'supercategory': c,
                'id': classes.index(c),
                'name': c
            }
        )
    with open(jsonPath, 'w') as j:
        json.dump(groudTruthDict, j, indent=2)
    # with open('val.json', 'w') as v:
    #     json.dump(resList, v, indent=2)

def get_coco_res(resPath, txtPath, imgRoot):
    with open(txtPath, 'r') as f:
        lines = f.readlines()
        for i, line in enumerate(tqdm(lines)):
            imgPath = os.path.join(imgRoot, line.strip()+'.jpg')
            assert os.path.exists(imgPath)



if __name__ == '__main__':
    txtPath = os.path.join(txtPath, 'trainval.txt')
    assert os.path.exists(txtPath) == True
    convert(txtPath=txtPath, annRoot=annPath)



